#!/usr/bin/python

from sqlalchemy import create_engine, Column, Integer, Boolean, String
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import json

engine = create_engine("sqlite:///testdb.db")
Session = sessionmaker(bind=engine)
Base = declarative_base()

class Task(Base):
    """A hash to be broken"""
    __tablename__ = "tasks"

    id = Column(Integer, primary_key=True)
    priority = Column(Integer)
    complete = Column(Boolean)
    all_assigned = Column(Boolean)
    salt = Column(String)
    hash_type = Column(String)
    hash = Column(String)
    plaintext = Column(String)
    min_length = Column(Integer)
    max_length = Column(Integer)
    markov_lowest = Column(Integer)
    markov_highest = Column(Integer)
    markov_possibles = Column(Integer)


class Client(Base):
    """A client that will be requesting workunits"""
    __tablename__ = "clients"

    id = Column(Integer, primary_key=True)
    hashes_per_second = Column(Integer)
    current_workunit_id = Column(Integer)

class Workunit(Base):
    """A list of potential solutions or instructions to generate such a list
       using incremental or Markov, and the hash to evaluate them against"""
    __tablename__ = "workunits"

    id = Column(Integer, primary_key=True)
    task_id = Column(Integer)
    client_id = Column(Integer)
    complete = Column(Boolean)
    markov_start = Column(Integer)
    markov_end = Column(Integer)

    def to_client(self):
        """Generate the data to be sent to the client for this workunit"""
        # Get the parent task from the database
        self.session  = Session()
        task = session.query(Task).filter_by(id=self.task_id).first()
        
        # Construct the data
        data = { "hash" : task.hash, }
        if task.markov_highest > 0:
            data["markov_highest"] = task.markov_highest
            data["markov_lowest"] = task.markov_lowest
            data["markov_start"] = self.markov_start
            data["markov_end"] = self.markov_end
        if task.max_length > 0:
            data["max_length"] = task.max_length
            data["min_length"] = task.min_length
        if task.salt != "":
            data["salt"] = task.salt
        if task.hash_type != "":
            data["hash_type"] = task.hash_type

        return json.dumps(data)
